Skip to content

[fix] mismatch dim during capture graph if with --gpu-memory-utilization - #40719

Closed
ir1ka wants to merge 2 commits into
vllm-project:mainfrom
ir1ka:tensor-dim-mismatch
Closed

[fix] mismatch dim during capture graph if with --gpu-memory-utilization#40719
ir1ka wants to merge 2 commits into
vllm-project:mainfrom
ir1ka:tensor-dim-mismatch

Conversation

@ir1ka

@ir1ka ir1ka commented Apr 23, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Purpose

Fix #40716.

Test Plan

Test model nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 on 4x RTX 3090 with pr #38985 based on version 0.19.2rc1.dev134+gfe9c3d6c5.

Although untested, this might also work for #39202.

Test Result

Start the vllm serve successful, and crash with error RuntimeError: The size of tensor a (58) must match the size of tensor b (63) at non-singleton dimension 1 if w/o this pr.

The test command:

VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=1 VLLM_MARLIN_USE_ATOMIC_ADD=1 VLLM_FLOAT32_MATMUL_PRECISION=high PYTORCH_ALLOC_CONF=expandable_segments:True CUDA_DEVICE_ORDER=PCI_BUS_ID vllm serve /models/nv-community/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 --served-model-name Nemotron-3-Super --tensor-parallel-size 4 --enable-expert-parallel --trust-remote-code --enable-auto-tool-choice --tool-call-parser qwen3_coder --reasoning-parser nemotron_v3 --gpu-memory-utilization 0.95 --max-model-len auto --async-scheduling --enable-prefix-caching --enable-chunked-prefill --max-num-seqs 4

Results:

$ docker compose exec lm-eval.ampere lm_eval --model local-completions  --model_args base_url=http://127.0.0.1:8000/v1/completions,model=Nemotron-3-Super,tokenizer=/models/nv-community/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 --trust_remote_code --tasks lambada_openai --batch_size auto

local-completions ({'base_url': 'http://127.0.0.1:8000/v1/completions', 'model': 'Nemotron-3-Super', 'tokenizer': '/models/nv-community/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4'}), gen_kwargs: ({}), limit: None, num_fewshot: None, batch_size: auto
|    Tasks     |Version|Filter|n-shot|  Metric  |   |Value |   |Stderr|
|--------------|------:|------|-----:|----------|---|-----:|---|-----:|
|lambada_openai|      1|none  |     0|acc       |↑  |0.7603|±  |0.0059|
|              |       |none  |     0|perplexity|↓  |3.0028|±  |0.0604|
Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

@ir1ka
ir1ka requested a review from njhill as a code owner April 23, 2026 17:21
claude[bot]
claude Bot reviewed Apr 23, 2026
View reviewed changes

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@mergify mergify Bot added the v1 label Apr 23, 2026
gemini-code-assist[bot]
gemini-code-assist Bot reviewed Apr 23, 2026
View reviewed changes

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request updates the GPUModelRunner to track and handle the re-initialization of InputBatch based on max_num_blocks and max_model_len. However, the may_reinitialize_input_batch method is missing the cp_kv_cache_interleave_size parameter when instantiating InputBatch, which could cause Context Parallelism to malfunction by reverting to default values.

Comment thread vllm/v1/worker/gpu_model_runner.py
@ir1ka
ir1ka force-pushed the tensor-dim-mismatch branch 3 times, most recently from c170809 to 744859b Compare April 30, 2026 16:24
@ir1ka
ir1ka force-pushed the tensor-dim-mismatch branch 2 times, most recently from 8142a3f to 22c3832 Compare May 4, 2026 09:19
@ir1ka
ir1ka force-pushed the tensor-dim-mismatch branch 2 times, most recently from 31d4ccd to 282472b Compare May 17, 2026 15:10
@ir1ka
ir1ka force-pushed the tensor-dim-mismatch branch from 282472b to 32c5cd4 Compare June 5, 2026 20:27
@ir1ka
ir1ka force-pushed the tensor-dim-mismatch branch from 32c5cd4 to c81cf71 Compare June 14, 2026 01:33
ir1ka added 2 commits June 24, 2026 01:20
@ir1ka
[fix] mismatch dim during capture graph if with --gpu-memory-utilization
a1df8b9 
Signed-off-by: IriKa Qiu <qiujie.jq@gmail.com>
@ir1ka
Add cp_kv_cache_interleave_size in may_reinitialize_input_batch to av…
9b1997d 
…oid reset to 1

Signed-off-by: IriKa Qiu <qiujie.jq@gmail.com>
@ir1ka
ir1ka force-pushed the tensor-dim-mismatch branch from c81cf71 to 9b1997d Compare June 23, 2026 17:20
@mergify

mergify Bot commented Jul 10, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @ir1ka.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Jul 10, 2026
@njhill njhill added the mrv1-only Issues/PRs which apply only to Model Runner V1 (not applicable to Model Runner V2) label Aug 28, 2026
@ir1ka ir1ka closed this Sep 11, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@claude claude[bot] claude[bot] left review comments
@njhill njhill Awaiting requested review from njhill njhill is a code owner
+1 more reviewer
@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

mrv1-only Issues/PRs which apply only to Model Runner V1 (not applicable to Model Runner V2) needs-rebase v1

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[Bug]: The size of tensor a (34) must match the size of tensor b (63) at non-singleton dimension 1

2 participants

@ir1ka @njhill